import sys
import pandas as pd
import numpy as np
from scipy import optimize
import matplotlib.pyplot as plt
import seaborn as sns
from matplotlib import rcParams

# 设置日志文件
LOG_FILE = 'A/prediction_log.txt'
log_file = open(LOG_FILE, 'w', encoding='utf-8')
original_stdout = sys.stdout
sys.stdout = log_file

# 读取四类文物统计表
def load_statistics(file_path):
    """读取四类文物统计表并整理为结构化字典"""
    sheets = {
        '铅钡无风化': '铅钡-无风化',
        '铅钡风化': '铅钡-风化',
        '高钾无风化': '高钾-无风化',
        '高钾风化': '高钾-风化'
    }
    
    stats_dict = {}
    
    for sheet_name, category in sheets.items():
        df = pd.read_excel(file_path, sheet_name=sheet_name)
        
        # 创建成分字典
        comp_dict = {}
        for _, row in df.iterrows():
            comp_name = row['成分']
            comp_dict[comp_name] = {
                'mean': row['均值'],
                'variance': row['方差'],
                'cv': row['变异系数'],
                'skew': row['偏度'],
                'kurt': row['峰度']
            }
        
        # 添加到统计字典
        glass_type, weather_status = category.split('-')
        if glass_type not in stats_dict:
            stats_dict[glass_type] = {}
        stats_dict[glass_type][weather_status] = comp_dict
    
    return stats_dict

# 玻璃类型分类函数
def classify_glass_type(row):
    """根据文物类型列判断玻璃类型"""
    return row['类型']

# 风化反演模型核心函数
def weathering_inversion(weathered_comp, glass_type, stats):
    """物理化学机制反演核心函数"""
    # 获取统计参数
    u_stats = stats[glass_type]['无风化']
    w_stats = stats[glass_type]['风化']
    
    results = {}
    uncertainties = {}
    
    # 对每个成分进行反演
    for comp, w_value in weathered_comp.items():
        # 跳过缺失或零值成分
        if comp not in u_stats or w_value == 0:
            results[comp] = w_value
            uncertainties[comp] = 0
            continue
            
        # 1. 计算基本反演值（基于均值比例）
        u_mean = u_stats[comp]['mean']
        w_mean = w_stats[comp]['mean']
        
        # 避免除零错误
        ratio = u_mean / w_mean if w_mean != 0 else 1.0
        base_pred = w_value * ratio
        
        # 2. 应用物理化学约束
        if comp == '二氧化硅(SiO2)':
            # SiO₂流失不超过80%
            max_loss = 0.8 * base_pred
            base_pred = max(base_pred, w_value + max_loss)
            
        elif comp == '二氧化硫(SO2)':
            # SO₂增加不超过20%
            max_gain = 0.2 * base_pred
            base_pred = min(base_pred, w_value - max_gain)
            
        elif glass_type == '铅钡' and comp == '氧化铅(PbO)':
            # 铅钡玻璃特征成分最低含量约束
            base_pred = max(base_pred, 15)
            
        elif glass_type == '铅钡' and comp == '氧化钡(BaO)':
            # 铅钡玻璃特征成分最低含量约束
            base_pred = max(base_pred, 5)
            
        elif glass_type == '高钾' and comp == '氧化钾(K2O)':
            # 高钾玻璃特征成分最低含量约束
            base_pred = max(base_pred, 8)
            
        # 3. 贝叶斯优化调整
        def objective(x):
            """目标函数：最小化与统计分布的差异"""
            # 先验差异（与未风化均值）
            prior_diff = abs(x - u_mean) / u_mean if u_mean != 0 else 0
            
            # 风化模型差异
            model_diff = abs(w_value - (x * w_mean / u_mean)) if u_mean != 0 and w_mean != 0 else 0
            
            # 分布特性差异（变异系数）
            cv_diff = abs(u_stats[comp]['cv'] - w_stats[comp]['cv'])
            
            return prior_diff + 0.5*model_diff + 0.3*cv_diff
        
        # 优化求解范围
        lower_bound = max(0, base_pred * 0.7)
        upper_bound = min(100, base_pred * 1.3)
        
        # 执行优化
        try:
            optimized = optimize.minimize_scalar(
                objective, 
                bounds=(lower_bound, upper_bound),
                method='bounded'
            )
            final_pred = optimized.x
        except Exception as e:
            print(f"优化失败 ({comp}): {e}, 使用基础预测值")
            final_pred = base_pred
        
        results[comp] = final_pred
        uncertainties[comp] = np.sqrt(u_stats[comp]['variance'])
    
    return results, uncertainties

# 结果后处理与验证
def post_process(prediction, uncertainties, glass_type, stats):
    """后处理确保结果合理性"""
    # 1. 非负约束
    for comp in prediction:
        prediction[comp] = max(0, prediction[comp])
    
    # 2. 总和约束（归一化到100%）
    total = sum(prediction.values())
    if total < 85 or total > 105:
        scale_factor = 100 / total
        for comp in prediction:
            prediction[comp] *= scale_factor
            uncertainties[comp] *= scale_factor
    
    # 3. 添加可靠性评估
    reliability_score = 5.0  # 最高5分
    for comp, value in prediction.items():
        if comp in stats[glass_type]['无风化']:
            u_mean = stats[glass_type]['无风化'][comp]['mean']
            u_std = np.sqrt(stats[glass_type]['无风化'][comp]['variance'])
            deviation = abs(value - u_mean) / u_std if u_std > 0 else 0
            # 每个成分的偏差影响可靠性
            reliability_score -= min(deviation, 1.0) * 0.3  # 最大每个成分扣0.3分
    
    reliability_score = max(1.0, min(5.0, reliability_score))  # 保持在1-5分
    
    return prediction, uncertainties, reliability_score

# 完整预测流程
def predict_original_composition(weathered_composition, glass_type, stats):
    """完整预测流程"""
    # 1. 应用反演模型
    raw_prediction, uncertainties = weathering_inversion(weathered_composition, glass_type, stats)
    
    # 2. 后处理
    final_prediction, final_uncertainties, reliability = post_process(
        raw_prediction, uncertainties, glass_type, stats
    )
    
    return {
        'prediction': final_prediction,
        'uncertainty': final_uncertainties,
        'reliability': reliability
    }

# 主函数
def main():
    print("="*50)
    print("古代玻璃文物风化前成分预测系统")
    print("="*50)
    
    # 1. 加载统计表
    stats_data = load_statistics("handled/四类文物统计.xlsx")
    print("成功加载四类文物统计表")
    
    # 2. 加载表单2数据（预处理后的数据）
    df = pd.read_excel("handled/预处理.xlsx")
    print(f"成功加载预处理数据，共 {len(df)} 条记录")
    
    # 3. 筛选风化样本
    # 注意：这里使用"表面风化"列，值为"风化"的样本
    weathered_df = df[df['表面风化'] == '风化'].copy()
    print(f"共筛选出 {len(weathered_df)} 个风化样本")
    
    # 4. 提取化学成分列名
    chem_columns = [
        '二氧化硅(SiO2)', '氧化钠(Na2O)', '氧化钾(K2O)', '氧化钙(CaO)', '氧化镁(MgO)',
        '氧化铝(Al2O3)', '氧化铁(Fe2O3)', '氧化铜(CuO)', '氧化铅(PbO)', '氧化钡(BaO)',
        '五氧化二磷(P2O5)', '氧化锶(SrO)', '氧化锡(SnO2)', '二氧化硫(SO2)'
    ]
    
    # 5. 执行预测
    results = []
    for idx, row in weathered_df.iterrows():
        sample_id = row['文物采样点']
        glass_type = classify_glass_type(row)
        
        # 提取化学成分数据
        chem_data = {col: row[col] for col in chem_columns if col in row}
        
        print(f"\n开始预测样本 {sample_id} ({glass_type}玻璃)...")
        
        # 执行预测
        prediction_result = predict_original_composition(chem_data, glass_type, stats_data)
        
        # 收集结果
        for chem, value in prediction_result['prediction'].items():
            unc = prediction_result['uncertainty'].get(chem, 0)
            results.append({
                '文物采样点': sample_id,
                '文物编号': row['文物编号'],
                '玻璃类型': glass_type,
                '成分': chem,
                '风化后含量(%)': row[chem] if chem in row else 0,
                '预测原始含量(%)': value,
                '不确定性(%)': unc,
                '可靠性评分': prediction_result['reliability']
            })
        
        # 打印简要结果
        print(f"可靠性评分: {prediction_result['reliability']:.1f}/5.0")
        print(f"主要成分预测:")
        print(f"  SiO₂: {prediction_result['prediction']['二氧化硅(SiO2)']:.1f}% "
              f"(风化后: {row['二氧化硅(SiO2)']:.1f}%)")
        if glass_type == '铅钡':
            print(f"  PbO: {prediction_result['prediction']['氧化铅(PbO)']:.1f}% "
                  f"(风化后: {row['氧化铅(PbO)']:.1f}%)")
            print(f"  BaO: {prediction_result['prediction']['氧化钡(BaO)']:.1f}% "
                  f"(风化后: {row['氧化钡(BaO)']:.1f}%)")
        else:
            print(f"  K₂O: {prediction_result['prediction']['氧化钾(K2O)']:.1f}% "
                  f"(风化后: {row['氧化钾(K2O)']:.1f}%)")
    
    # 6. 保存结果到Excel
    results_df = pd.DataFrame(results)
    output_file = "A/所有风化文物原始成分预测.xlsx"
    results_df.to_excel(output_file, index=False)
    print(f"\n预测结果已保存到: {output_file}")
    
   # 7. 添加可视化分析
    def create_visualization(results_df, weathered_df):
        """创建风化前后成分变化的可视化图表"""
        # 设置中文字体
        plt.rcParams['font.sans-serif'] = ['SimHei', 'Arial Unicode MS']
        plt.rcParams['axes.unicode_minus'] = False
        
        # 准备数据
        comparison_data = []
        for _, row in results_df.iterrows():
            sample_id = row['文物采样点']
            component = row['成分']
            weathered_content = row['风化后含量(%)']
            predicted_content = row['预测原始含量(%)']
            
            comparison_data.append({
                '文物采样点': sample_id,
                '成分': component,
                '风化后含量': weathered_content,
                '预测原始含量': predicted_content,
                '变化量': predicted_content - weathered_content,
                '变化率(%)': ((predicted_content - weathered_content) / predicted_content * 100) if predicted_content != 0 else 0
            })
        
        comp_df = pd.DataFrame(comparison_data)
        
        # 1. 绘制所有成分的平均变化柱状图
        plt.figure(figsize=(15, 8))
        avg_change = comp_df.groupby('成分')[['风化后含量', '预测原始含量']].mean()
        
        x = range(len(avg_change.index))
        width = 0.35
        
        plt.bar([i - width/2 for i in x], avg_change['风化后含量'], width, 
                label='风化后平均含量', alpha=0.8, color='orange')
        plt.bar([i + width/2 for i in x], avg_change['预测原始含量'], width, 
                label='预测原始平均含量', alpha=0.8, color='skyblue')
        
        plt.xlabel('化学成分')
        plt.ylabel('含量 (%)')
        plt.title('所有风化文物各成分风化前后平均含量对比')
        plt.xticks(x, avg_change.index, rotation=45, ha='right')
        plt.legend()
        plt.grid(True, alpha=0.3)
        plt.tight_layout()
        plt.savefig('A/成分平均含量对比柱状图.png', dpi=300, bbox_inches='tight')
        plt.show()
        
        # 2. 绘制主要成分的变化率柱状图
        plt.figure(figsize=(12, 6))
        main_components = ['二氧化硅(SiO2)', '氧化钠(Na2O)', '氧化钾(K2O)', 
                          '氧化钙(CaO)', '氧化铅(PbO)', '氧化钡(BaO)']
        
        main_comp_data = comp_df[comp_df['成分'].isin(main_components)]
        avg_change_rate = main_comp_data.groupby('成分')['变化率(%)'].mean()
        
        colors = ['red' if x < 0 else 'green' for x in avg_change_rate.values]
        bars = plt.bar(range(len(avg_change_rate)), avg_change_rate.values, color=colors, alpha=0.7)
        
        plt.xlabel('主要化学成分')
        plt.ylabel('平均变化率 (%)')
        plt.title('主要成分风化过程中的平均变化率\n(负值表示流失，正值表示增加)')
        plt.xticks(range(len(avg_change_rate)), avg_change_rate.index, rotation=45, ha='right')
        plt.axhline(y=0, color='black', linestyle='-', alpha=0.3)
        plt.grid(True, alpha=0.3)
        
        # 添加数值标签
        for bar, value in zip(bars, avg_change_rate.values):
            plt.text(bar.get_x() + bar.get_width()/2, bar.get_height() + (1 if value > 0 else -3), 
                    f'{value:.1f}%', ha='center', va='bottom' if value > 0 else 'top')
        
        plt.tight_layout()
        plt.savefig('A/主要成分变化率柱状图.png', dpi=300, bbox_inches='tight')
        plt.show()
        
        # 3. 绘制各样本主要成分的折线图
        plt.figure(figsize=(16, 10))
        
        # 选择几个代表性样本
        unique_samples = comp_df['文物采样点'].unique()[:8]  # 取前8个样本避免图表过于拥挤
        
        for i, component in enumerate(main_components):
            plt.subplot(2, 3, i+1)
            
            comp_data = comp_df[comp_df['成分'] == component]
            sample_data = comp_data[comp_data['文物采样点'].isin(unique_samples)]
            
            x_pos = range(len(sample_data))
            plt.plot(x_pos, sample_data['风化后含量'], 'o-', label='风化后', 
                    color='orange', linewidth=2, markersize=6)
            plt.plot(x_pos, sample_data['预测原始含量'], 's-', label='预测原始', 
                    color='skyblue', linewidth=2, markersize=6)
            
            plt.title(f'{component}')
            plt.xlabel('样本编号')
            plt.ylabel('含量 (%)')
            plt.xticks(x_pos, [s.split('-')[-1] if '-' in s else s for s in sample_data['文物采样点']], 
                      rotation=45, ha='right')
            plt.legend(fontsize=8)
            plt.grid(True, alpha=0.3)
        
        plt.suptitle('代表性样本主要成分风化前后含量变化折线图', fontsize=14, y=0.98)
        plt.tight_layout()
        plt.savefig('A/样本成分变化折线图.png', dpi=300, bbox_inches='tight')
        plt.show()
        
        # 4. 绘制成分变化热力图
        plt.figure(figsize=(14, 8))
        
        # 创建变化率矩阵
        pivot_data = comp_df.pivot(index='成分', columns='文物采样点', values='变化率(%)')
        
        # 只显示前10个样本以避免过于拥挤
        pivot_data_subset = pivot_data.iloc[:, :10]
        
        sns.heatmap(pivot_data_subset, annot=True, fmt='.1f', cmap='RdYlBu_r', 
                   center=0, cbar_kws={'label': '变化率 (%)'})
        plt.title('各文物成分风化变化率热力图\n(红色表示流失，蓝色表示增加)')
        plt.xlabel('文物采样点')
        plt.ylabel('化学成分')
        plt.xticks(rotation=45, ha='right')
        plt.yticks(rotation=0)
        plt.tight_layout()
        plt.savefig('A/成分变化热力图.png', dpi=300, bbox_inches='tight')
        plt.show()
        
        # 5. 绘制玻璃类型分组的成分变化对比
        plt.figure(figsize=(15, 6))
        
        # 合并玻璃类型信息
        glass_type_map = dict(zip(weathered_df['文物采样点'], weathered_df['类型']))
        comp_df['玻璃类型'] = comp_df['文物采样点'].map(glass_type_map)
        
        # 按玻璃类型分组的主要成分平均变化率
        type_comp_change = comp_df[comp_df['成分'].isin(main_components)].groupby(['玻璃类型', '成分'])['变化率(%)'].mean().unstack()
        
        type_comp_change.plot(kind='bar', figsize=(12, 6), width=0.8)
        plt.title('不同玻璃类型主要成分的平均变化率对比')
        plt.xlabel('玻璃类型')
        plt.ylabel('平均变化率 (%)')
        plt.xticks(rotation=0)
        plt.legend(title='化学成分', bbox_to_anchor=(1.05, 1), loc='upper left')
        plt.axhline(y=0, color='black', linestyle='-', alpha=0.3)
        plt.grid(True, alpha=0.3)
        plt.tight_layout()
        plt.savefig('A/玻璃类型成分变化对比.png', dpi=300, bbox_inches='tight')
        plt.show()
        
        print("\n📊 可视化图表已生成并保存到 A/ 目录:")
        print("  - 成分平均含量对比柱状图.png")
        print("  - 主要成分变化率柱状图.png") 
        print("  - 样本成分变化折线图.png")
        print("  - 成分变化热力图.png")
        print("  - 玻璃类型成分变化对比.png")
    
    # 调用可视化函数
    print("\n🎨 开始生成可视化图表...")
    create_visualization(results_df, weathered_df)
    
    # 8. 关闭日志文件
    sys.stdout = original_stdout
    log_file.close()
    print(f"日志已保存到: {LOG_FILE}")

if __name__ == "__main__":
    main()